/*
 * Wazuh storeRemediationsModel
 * Copyright (C) 2015, Wazuh Inc.
 * May 03, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "updateHotfixes_test.hpp"
#include "flatbuffers/idl.h"
#include "flatbuffers/util.h"
#include <memory>

namespace NSUpdateHotfixesTest
{
    constexpr auto COMMON_DATABASE_DIR {"queue/vd"}; //<<Used for all databases
    const std::string FLATBUFFER_SCHEMA {FLATBUFFER_SCHEMAS_DIR "/cve5.fbs"};
    const char* FB_INCLUDE_DIRECTORIES[] = {FLATBUFFER_SCHEMAS_DIR, nullptr};

    const auto JSON_CVE5_VALID_ONE_BLOCK = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-1111-1111",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

    const auto JSON_CVE5_VALID_MULTIPLE_BLOCKS = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-2222-2222",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-4000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-5000","KBT-6000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-7000","KBT-8000","KBT-9000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

    const auto JSON_CVE5_VALID_MULTIPLE_BLOCKS_SOME_WITH_NO_ANYOF = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-3333-3333",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "anyOf":["KBT-800","KBT-1000","KBT-3000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-4000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        },
                        {
                            "anyOf":["KBT-7000","KBT-8000","KBT-9000"],
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

    const auto JSON_CVE5_EMPTY_UPDATES = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-4444-4444",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                    "windows": [
                        {
                            "products": ["Windows 10 - HastaLaVistaBaby", "Windows 11 - RiseOfTheMachines", "Windows 12 - Genisys"],
                            "type": "update"
                        }
                    ]
                }
            }
        }
    }
)";

    const auto JSON_CVE5_NO_WINDOWS = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-4444-4444",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
                "x_remediations": {
                }
            }
        }
    }
)";

    const auto JSON_CVE5_NO_REMEDIATIONS = R"(
    {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.0",
        "cveMetadata": {
            "cveId": "CVE-5555-5555",
            "assignerOrgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6",
            "state": "PUBLISHED"
        },
        "containers": {
            "cna": {
                "providerMetadata": {
                    "orgId": "b3476cb9-2e3d-41a6-98d0-0f47421a65b6"
                },
                "problemTypes": [
                    {
                        "descriptions": [
                            {
                                "lang": "en",
                                "description": "CWE-78 OS Command Injection"
                            }
                        ]
                    }
                ],
                "affected": [
                    {
                        "vendor": "Example.org",
                        "product": "Example Enterprise",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "lessThan": "1.0.6",
                                "versionType": "semver"
                            }
                        ],
                        "defaultStatus": "unaffected"
                    }
                ],
                "descriptions": [
                    {
                        "lang": "en",
                        "value": "OS Command Injection vulnerability parseFilename function of example.php in the Web Management Interface of Example.org Example Enterprise on Windows, MacOS and XT-4500 allows remote unauthenticated attackers to escalate privileges.\n\nThis issue affects:\n  *  1.0 versions before 1.0.6\n  *  2.1 versions from 2.16 until 2.1.9."
                    }
                ],
                "references": [
                    {
                        "url": "https://example.org/ESA-22-11-CVE-1337-1234"
                    }
                ],
            }
        }
    }
)";
}; // namespace NSUpdateHotfixesTest

using namespace NSUpdateHotfixesTest;

void UpdateHotfixesTest::SetUp()
{
    std::filesystem::remove_all(COMMON_DATABASE_DIR);
    std::filesystem::create_directories(COMMON_DATABASE_DIR);
};

void UpdateHotfixesTest::TearDown()
{
    std::filesystem::remove_all(COMMON_DATABASE_DIR);
};

TEST_F(UpdateHotfixesTest, UpdatesWindowsRemediationMultipleBlocks)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

    const auto cveId = entry->cveMetadata()->cveId()->str();

    // Validate the data was inserted
    std::vector<std::string> keys {"KBT-800_" + cveId,
                                   "KBT-1000_" + cveId,
                                   "KBT-3000_" + cveId,
                                   "KBT-4000_" + cveId,
                                   "KBT-5000_" + cveId,
                                   "KBT-6000_" + cveId,
                                   "KBT-7000_" + cveId,
                                   "KBT-8000_" + cveId,
                                   "KBT-9000_" + cveId};

    for (const auto& key : keys)
    {
        rocksdb::PinnableSlice slice;
        EXPECT_TRUE(rocksDBWrapper->get(key, slice, HOTFIXES_APPLICATIONS_COLUMN)) << "Unable to find key: " << key;
    }
}

TEST_F(UpdateHotfixesTest, SkipsEmptyRemediations)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_NO_REMEDIATIONS);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

    // The column should not have been created
    EXPECT_FALSE(rocksDBWrapper->columnExists(HOTFIXES_APPLICATIONS_COLUMN));
}

TEST_F(UpdateHotfixesTest, SkipsNoWindows)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_NO_WINDOWS);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

    // The column should not have been created
    EXPECT_FALSE(rocksDBWrapper->columnExists(HOTFIXES_APPLICATIONS_COLUMN));
}

TEST_F(UpdateHotfixesTest, SkipsEmptyUpdates)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_EMPTY_UPDATES);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

    // Check that no data was inserted
    auto counter = 0;
    for (auto [key, value] : rocksDBWrapper->seek("KBT", HOTFIXES_APPLICATIONS_COLUMN))
    {
        counter++;
    }
    EXPECT_EQ(counter, 0);
}

TEST_F(UpdateHotfixesTest, MultipleVulnerabilities)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Expected keys to be inserted
    std::set<std::string> expectedKeys {};

    {
        // Parse the test data
        std::string schemaStr;
        flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

        flatbuffers::Parser parser;
        parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
        parser.Parse(JSON_CVE5_VALID_ONE_BLOCK);

        // Create and verify the Entry object.
        auto buffer = parser.builder_.GetBufferPointer();
        flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
        cve_v5::VerifyEntryBuffer(verifier);

        const auto entry = cve_v5::GetEntry(buffer);

        // Insert the vulnerability
        EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

        // Add the key to the list
        const auto cveId = entry->cveMetadata()->cveId()->str();
        expectedKeys.insert("KBT-800_" + cveId);
        expectedKeys.insert("KBT-1000_" + cveId);
        expectedKeys.insert("KBT-3000_" + cveId);
    }

    {
        // Parse the test data
        std::string schemaStr;
        flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

        flatbuffers::Parser parser;
        parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
        parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS);

        // Create and verify the Entry object.
        auto buffer = parser.builder_.GetBufferPointer();
        flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
        cve_v5::VerifyEntryBuffer(verifier);

        const auto entry = cve_v5::GetEntry(buffer);

        // Insert the vulnerability
        EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

        // Add the key to the list
        const auto cveId = entry->cveMetadata()->cveId()->str();
        expectedKeys.insert("KBT-4000_" + cveId);
        expectedKeys.insert("KBT-5000_" + cveId);
        expectedKeys.insert("KBT-6000_" + cveId);
        expectedKeys.insert("KBT-7000_" + cveId);
        expectedKeys.insert("KBT-8000_" + cveId);
        expectedKeys.insert("KBT-9000_" + cveId);
    }

    rocksdb::PinnableSlice slice;
    for (const auto& key : expectedKeys)
    {
        EXPECT_TRUE(rocksDBWrapper->get(key, slice, HOTFIXES_APPLICATIONS_COLUMN)) << "Unable to find key: " << key;
    }
}

TEST_F(UpdateHotfixesTest, RemoveSkipNoDatabase)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_EMPTY_UPDATES);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::removeHotfix(entry, rocksDBWrapper.get()));
}

TEST_F(UpdateHotfixesTest, RemoveSkipNoWindows)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Create the column
    EXPECT_NO_THROW(rocksDBWrapper->createColumn(HOTFIXES_APPLICATIONS_COLUMN));

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_NO_WINDOWS);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::removeHotfix(entry, rocksDBWrapper.get()));
}

TEST_F(UpdateHotfixesTest, RemoveSkipNoRemediations)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Create the column
    EXPECT_NO_THROW(rocksDBWrapper->createColumn(HOTFIXES_APPLICATIONS_COLUMN));

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_NO_REMEDIATIONS);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::removeHotfix(entry, rocksDBWrapper.get()));
}

TEST_F(UpdateHotfixesTest, RemoveSkipEmptyUpdates)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Create the column
    EXPECT_NO_THROW(rocksDBWrapper->createColumn(HOTFIXES_APPLICATIONS_COLUMN));

    // Parse the test data
    std::string schemaStr;
    flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

    flatbuffers::Parser parser;
    parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
    parser.Parse(JSON_CVE5_EMPTY_UPDATES);

    // Create and verify the Entry object.
    auto buffer = parser.builder_.GetBufferPointer();
    flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
    cve_v5::VerifyEntryBuffer(verifier);

    const auto entry = cve_v5::GetEntry(buffer);

    // Call the updateRemediation function with the test data
    EXPECT_NO_THROW(UpdateHotfixes::removeHotfix(entry, rocksDBWrapper.get()));
}

TEST_F(UpdateHotfixesTest, StoreAndRemove)
{
    // Create a mock RocksDBWrapper object
    std::unique_ptr<Utils::IRocksDBWrapper> rocksDBWrapper = std::make_unique<Utils::RocksDBWrapper>(DATABASE_PATH);

    // Expected and non expected keys
    std::unordered_set<std::string> expectedKeys {};
    std::unordered_set<std::string> nonExpectedKeys {};

    {
        // Parse the test data
        std::string schemaStr;
        flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

        flatbuffers::Parser parser;
        parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
        parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS);

        // Create and verify the Entry object.
        auto buffer = parser.builder_.GetBufferPointer();
        flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
        cve_v5::VerifyEntryBuffer(verifier);

        const auto entry = cve_v5::GetEntry(buffer);

        // Insert the vulnerability
        EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

        // Update the expected keys
        const auto cveId = entry->cveMetadata()->cveId()->str();
        expectedKeys.insert("KBT-800_" + cveId);
        expectedKeys.insert("KBT-1000_" + cveId);
        expectedKeys.insert("KBT-3000_" + cveId);
        expectedKeys.insert("KBT-4000_" + cveId);
        expectedKeys.insert("KBT-5000_" + cveId);
        expectedKeys.insert("KBT-6000_" + cveId);
        expectedKeys.insert("KBT-7000_" + cveId);
        expectedKeys.insert("KBT-8000_" + cveId);
        expectedKeys.insert("KBT-9000_" + cveId);
    }

    {
        // Parse the test data
        std::string schemaStr;
        flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

        flatbuffers::Parser parser;
        parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
        parser.Parse(JSON_CVE5_VALID_ONE_BLOCK);

        // Create and verify the Entry object.
        auto buffer = parser.builder_.GetBufferPointer();
        flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
        cve_v5::VerifyEntryBuffer(verifier);

        const auto entry = cve_v5::GetEntry(buffer);

        // Insert the vulnerability
        EXPECT_NO_THROW(UpdateHotfixes::storeVulnerabilityHotfixes(entry, rocksDBWrapper.get()));

        // Update the expected keys
        const auto cveId = entry->cveMetadata()->cveId()->str();
        expectedKeys.insert("KBT-800_" + cveId);
        expectedKeys.insert("KBT-1000_" + cveId);
        expectedKeys.insert("KBT-3000_" + cveId);
    }

    {
        // Parse the test data
        std::string schemaStr;
        flatbuffers::LoadFile(FLATBUFFER_SCHEMA.c_str(), false, &schemaStr);

        flatbuffers::Parser parser;
        parser.Parse(schemaStr.c_str(), FB_INCLUDE_DIRECTORIES);
        parser.Parse(JSON_CVE5_VALID_MULTIPLE_BLOCKS);

        // Create and verify the Entry object.
        auto buffer = parser.builder_.GetBufferPointer();
        flatbuffers::Verifier verifier(buffer, parser.builder_.GetSize());
        cve_v5::VerifyEntryBuffer(verifier);

        const auto entry = cve_v5::GetEntry(buffer);

        // Remove the hotfixes
        EXPECT_NO_THROW(UpdateHotfixes::removeHotfix(entry, rocksDBWrapper.get()));

        // Update the non expected keys
        const auto cveId = entry->cveMetadata()->cveId()->str();
        nonExpectedKeys.insert("KBT-800_" + cveId);
        nonExpectedKeys.insert("KBT-1000_" + cveId);
        nonExpectedKeys.insert("KBT-3000_" + cveId);
        nonExpectedKeys.insert("KBT-4000_" + cveId);
        nonExpectedKeys.insert("KBT-5000_" + cveId);
        nonExpectedKeys.insert("KBT-6000_" + cveId);
        nonExpectedKeys.insert("KBT-7000_" + cveId);
        nonExpectedKeys.insert("KBT-8000_" + cveId);
        nonExpectedKeys.insert("KBT-9000_" + cveId);

        // Update the expected keys
        expectedKeys.erase("KBT-800_" + cveId);
        expectedKeys.erase("KBT-1000_" + cveId);
        expectedKeys.erase("KBT-3000_" + cveId);
        expectedKeys.erase("KBT-4000_" + cveId);
        expectedKeys.erase("KBT-5000_" + cveId);
        expectedKeys.erase("KBT-6000_" + cveId);
        expectedKeys.erase("KBT-7000_" + cveId);
        expectedKeys.erase("KBT-8000_" + cveId);
        expectedKeys.erase("KBT-9000_" + cveId);
    }

    // Remaining hotfixes:
    // KBT-800, KBT-1000, KBT-3000
    rocksdb::PinnableSlice slice;

    for (const auto& key : expectedKeys)
    {
        EXPECT_TRUE(rocksDBWrapper->get(key, slice, HOTFIXES_APPLICATIONS_COLUMN)) << "Unable to find key: " << key;
    }

    for (const auto& key : nonExpectedKeys)
    {
        EXPECT_FALSE(rocksDBWrapper->get(key, slice, HOTFIXES_APPLICATIONS_COLUMN)) << "Found key: " << key;
    }
}
